
import glob
import random
import os
from PIL import Image, ImageEnhance,ImageFilter
import numpy as np
from mindspore import Tensor
from mindspore import dtype as mstype
import mindspore.dataset as ds

class ImageDataset:
    def __init__(self, root, mode='train'):

        self.files_A = sorted(glob.glob(os.path.join(root, '%s/A' % mode) + '/*.*'))
        self.files_B = sorted(glob.glob(os.path.join(root, '%s/B' % mode) + '/*.*'))
        self.files_Label = sorted(glob.glob(os.path.join(root, '%s/label' % mode) + '/*.*'))

    def __getitem__(self, index):

        input_A = Image.open(self.files_A[index % len(self.files_A)]).convert('RGB')
        input_B = Image.open(self.files_B[index % len(self.files_B)]).convert('RGB')
        label = Image.open(self.files_Label[index % len(self.files_Label)]).convert('L')

        imgA = np.asarray(input_A).astype(np.float32)
        imgB = np.asarray(input_B).astype(np.float32)
        label = np.asarray(label).astype(np.uint8)

        imgA = Tensor(imgA / 255.).transpose(-1, 0, 1)
        imgB = Tensor(imgB / 255.).transpose(-1, 0, 1)
        label = Tensor(label // 255)


        return imgA, imgB,  label

    def __len__(self):
        return len(self.files_A)

def create_dataset_new(batch = 16, mode='train'):
    # instanceof dataset
    dataset_generator = ImageDataset(root=r'D:\git_test\Levir_sample\LEVIR-train_val_test', mode=mode)

    dataset = ds.GeneratorDataset(dataset_generator, ["imgA", "imgB", "label"], shuffle=False)
    dataset = dataset.batch(batch_size=batch)
    return dataset


if __name__ == '__main__':
    dataset = ImageDataset(root=r'D:\git_test\Levir_sample\LEVIR-train_val_test')
    out = dataset.__getitem__(1)
    print(out)